import matplotlib.pyplot as plt
import numpy as np
import math

class Edge:
    """
    x, y : 边的端点
    length: 边长
    """
    def __init__(self, x, y, length):
        self.x = x
        self.y = y
        self.length = length

class UnionFindSet:
    def __init__(self, start, n):
        self.start = start 
        self.n = n
        self.pre = [0 for _ in range(self.n - self.start + 2)] 
        self.rank = [0 for _ in range(self.n - self.start + 2)]  

    def init(self):
        for i in range(self.start, self.n+1):
            self.pre[i] = i
            self.rank[i] = 1

    def find_pre(self, x):
        if self.pre[x] != x:
            self.pre[x] = self.find_pre(self.pre[x])
        return self.pre[x]

    def is_same(self, x, y):
        return self.find_pre(x) == self.find_pre(y)

    def unite(self, x, y):
        x = self.find_pre(x)
        y = self.find_pre(y)
        if x == y:
            return False
        if self.rank[x] > self.rank[y]:
            self.pre[y] = x
        else:
            if self.rank[x] == self.rank[y]:
                self.rank[y] += 1
            self.pre[x] = y
        return True

class Kruskal:
    """Kruskal算法 """
    def __init__(self, p, e):
        self.n = len(p) # n,m分别表示输入的点和边的个数
        self.m = len(e)
        self.e = e  # 存放录入的无向连通图的所有边
        self.p = p
        self.s = []  # 存放最小生成树里的所有边
        self.u = UnionFindSet(1, self.n)  

    def graphy(self):
        """ 所有边排序 """
        self.e.sort(key=lambda e: e.length)
        self.u.init()

    def run(self):
        """执行函数：求解录入的无向连通图的最小生成树 """
        for i in range(self.m):
            if self.u.unite(self.e[i].x, self.e[i].y):
                self.s.append(self.e[i])
            if len(self.s) == self.n - 1:
                break

    def draw(self):
        plt.figure(figsize=(6,6))
        plt.xlim(-5,105)     
        plt.ylim(-5,105)       
        plt.xlabel('X')    
        plt.ylabel('Y')    
        for i in range(self.n):
            plt.text(self.p[i][0] + 0.5,self.p[i][1] + 0.5, str(i),color='#0085c3')
            plt.plot(self.p[i][0], self.p[i][1],'o',color='#ff1493')
        for i in self.s:
            start = (self.p[i.x][0], self.p[i.y][0])
            end = (self.p[i.x][1], self.p[i.y][1])
            plt.plot(start, end, color='#ff1493')
        plt.title("Kruskal Algorithm - Minimum Spining Tree\n Union Find")
        plt.grid()
        plt.show()

def main():
    p, e = generate_graph(50) # 产生随机点
    kruskal = Kruskal(p, e)
    kruskal.graphy()
    kruskal.run()
    kruskal.print()
    kruskal.draw()


def generate_graph(Points_Num):
    NaivePoints = np.random.rand(Points_Num, 2)
    points = []
    edges = []
    for i in NaivePoints:
        points.append([i[0] * Points_Num * 2, i[1] * Points_Num * 2])
    for i in range(len(points) - 1):
        for j in range(i + 1, len(points)):
            x1, y1 = points[i][0], points[i][1]
            x2, y2 = points[j][0], points[j][1]
            edges.append(Edge(i, j, math.sqrt((x2 - x1)**2 + (y2 - y1)**2)))
    return points, edges

if __name__ == '__main__':
    main()





